package com.xiyuan.smartutils.client.http.ws;

import com.xiyuan.smartutils.constants.WebSocketConstants;

import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.util.Random;

/**
 * WebSocket 数据包
 * WebSocket数据包格式：<br>
 *
 * <a href="https://tools.ietf.org/html/rfc6455#section-5.2">RFC 6455</a>
 * <pre>
 *  0                   1                   2                   3
 *  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
 * +-+-+-+-+-------+-+-------------+-------------------------------+
 * |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
 * |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
 * |N|V|V|V|       |S|             |   (if payload len==126/127)   |
 * | |1|2|3|       |K|             |                               |
 * +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
 * |     Extended payload length continued, if payload len == 127  |
 * + - - - - - - - - - - - - - - - +-------------------------------+
 * |                               |Masking-key, if MASK set to 1  |
 * +-------------------------------+-------------------------------+
 * | Masking-key (continued)       |          Payload Data         |
 * +-------------------------------- - - - - - - - - - - - - - - - +
 * :                     Payload Data continued ...                :
 * + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
 * |                     Payload Data continued ...                |
 * +---------------------------------------------------------------+
 * @author xiyuan-lgz 2024-11-03 @version v1.0.0 新建与整理
 */
@SuppressWarnings("AlibabaLowerCamelCaseVariableNaming")
public class WebSocketPacket implements WebSocketConstants {
    private static final Random RANDOM = new Random();
    private byte[] data;
    private final PacketHeader header;
    
    public WebSocketPacket(Opcode opcode) {
        this((String) null, opcode);
    }
    
    public WebSocketPacket(String body) {
        this(body == null ? null : body.getBytes(), TEXT);
    }
    
    public WebSocketPacket(String body, Opcode opcode) {
        this(body == null ? null : body.getBytes(), opcode);
    }
    
    public WebSocketPacket(byte[] body) {
        this(body, new PacketHeader(BIN));
    }
    
    public WebSocketPacket(ByteBuffer body) {
        this(body, new PacketHeader(BIN));
    }
    
    public WebSocketPacket(ByteBuffer body, Opcode opcode) {
        this(body, new PacketHeader(opcode));
    }
    
    public WebSocketPacket(byte[] body, Opcode opcode) {
        this(body, new PacketHeader(opcode));
    }
    
    public WebSocketPacket(byte[] body, PacketHeader header) {
        int len = body == null ? 0 : body.length;
        if (header == null) {
            header = new PacketHeader(BIN, len, PayloadLength.of(len));
        }
        this.data = body;
        this.header = header.setPayloadLength(len);
    }
    
    public WebSocketPacket(ByteBuffer body, PacketHeader header) {
        int len = body == null ? 0 : body.remaining();
        if (header == null) {
            header = new PacketHeader(BIN, len, PayloadLength.of(len));
        }
        
        this.header = header;
        if (body != null) {
            byte[] buf = new byte[len];
            body.get(buf);
            this.data = buf;
        }
    }
    
    
    public byte[] getData() {
        return data;
    }
    
    public String getText() {
        return new String(data);
    }
    
    public PacketHeader getHeader() {
        return header;
    }
    
    public final String getText(Charset charset) {
        if (data == null) {
            return null;
        }
        return charset == null ? new String(data) : new String(data, charset);
    }
    
    /**
     * set mask
     */
    public void mask() {
        header.setMask(true);
    }
    
    public void unmask() {
        header.setMask(false);
    }
    
    public ByteBuffer toByteBuffer() {
        PayloadLength lenType = header.getPayloadLengthType();
        int len = header.getPayloadLength();
        
        ByteBuffer buf = ByteBuffer.allocate(2 + lenType.offset() + len + (header.isMask() ? 4 : 0));
        
        // FIN
        buf.put((byte) ((header.isFragmented() ? 0 : FIN_MASK) | header.getOpcode().intValue()));
        // payloadLen
        byte maskBit = (byte) (header.isMask() ? MASK_MASK : 0);
        switch (lenType) {
            case LEN_SHORT:
                buf.put((byte) (len | maskBit));
                break;
            case LEN_16:
                buf.put((byte) (lenType.byteValue() | maskBit));
                buf.putShort((short) len);
                break;
            case LEN_63:
                buf.put((byte) (lenType.byteValue() | maskBit));
                buf.putLong(len);
                break;
            default:
                break;
        }
        
        if (header.isMask()) {
            // 掩码加密处理
            byte[] mask_key = new byte[4];
            ByteBuffer.wrap(mask_key).putInt(RANDOM.nextInt());
            buf.put(mask_key);
            if (data != null) {
                for (int i = 0; i < data.length; i++) {buf.put((byte) (data[i] ^ mask_key[i % 4]));}
            }
        }
        else if (data != null) {
            buf.put(data);
        }
        
        buf.flip();
        return buf;
    }
    
    public int length() {
        return header.getPayloadLength();
    }
    
    
    public static class PacketHeader {
        
        private boolean fragmented;
        private boolean mask;
        private Opcode opcode;
        private int payloadLength;
        private PayloadLength payloadType;
        
        public PacketHeader(Opcode opcode) {
            this(opcode, 0, PayloadLength.LEN_SHORT);
        }
        
        public PacketHeader(Opcode opcode, int payloadLength, PayloadLength payloadType) {
            if (opcode == null) throw new IllegalArgumentException("Invalid opcode");
            if (payloadType == null) throw new IllegalArgumentException("Invalid payloadLength");
            this.opcode = opcode;
            this.fragmented = opcode == MID;
            this.payloadLength = payloadLength;
            this.payloadType = payloadType;
            
        }
        
        /**
         * is fragmented packet
         *
         * @return true|false
         */
        public boolean isFragmented() {
            return fragmented;
        }
        
        
        public Opcode getOpcode() {
            return opcode;
        }
        
        public PacketHeader setOpcode(Opcode opcode) {
            if (opcode == null) {
                return this;
            }
            this.opcode = opcode;
            this.fragmented = opcode == MID;
            return this;
        }
        
        public boolean isMask() {
            return mask;
        }
        
        public PacketHeader setMask(boolean mask) {
            this.mask = mask;
            return this;
        }
        
        public int getPayloadLength() {
            return payloadLength;
        }
        
        public PacketHeader setPayloadLength(int payloadLength) {
            this.payloadType = PayloadLength.of(payloadLength);
            this.payloadLength = payloadLength;
            return this;
        }
        
        public PayloadLength getPayloadLengthType() {
            return payloadType;
        }
    }
}
